#ifndef MEMORYMANAGEMENT_H
#define MEMORYMANAGEMENT_H

#include <Globals.h>
#include <Heap.h>
#include <BitArray.h>
#include <SerialDriver.h>

class PageDirectory;
class Page;

#define HEAP_START         0xC0000000
#define HEAP_INITIAL_SIZE  0x100000
#define HEAP_END           0xCFFFFFFF
#define USER_HEAP_START    0xD0000000
#define USER_HEAP_INITIAL_SIZE 0x100000
#define USER_HEAP_END      0xDFFFFFFF
#define STACK_START        (0xB0000000-0x4)
#define STACK_INITIAL_SIZE 0x10000
#define STACK_END          0xA0000000
#define KERNEL_START       0x100000

#define STACK_ADDRESS(x) ((u32int)x <= STACK_START && (u32int)x > STACK_END)
#define HEAP_ADDRESS(x)  ((u32int)x >= HEAP_START  && (u32int)x < USER_HEAP_START)
#define USER_HEAP_ADDRESS(x) ((u32int)x >= USER_HEAP_START && (u32int)x < 0xF0000000)

/**
  Handles all memory related events. Heap startup, allocation, deallocation,
  virtual memory etc.
**/
class MemoryManagement
{
public:

  /**
    Heap can call our private allocFrame() functions.
  **/
  friend class Heap;
  /**
    ResourceManager can access our userHeap member.
  **/
  friend class ResourceManager;
  
  /**
    Default constructor, called on bootup.
  **/
  MemoryManagement();

  ~MemoryManagement();

  /**
    Normal constructor - passes the address of end of memory. Initialises paging and sets up
    a standard kernel page directory. Enables paging, then maps some pages for the heap.
  **/
  void initialise(u32int memEnd);

  /**
    Allocate "size" bytes, returning the physical address of the segment allocated in
    physicalLocation if physicalLocation != NULL.
  **/
  void *malloc(u32int size, bool pageAlign=false, u32int *physicalLocation=NULL);
  
  /**
    Deallocate the memory allocated to p.
  **/
  void free(void *p);
  
  /**
    Allocate "size" bytes from the *user space heap*.
  **/
  void *umalloc(u32int size);
  
  /**
    Deallocate any memory allocated to p via umalloc.
  **/
  void ufree(void *p);
  
  /**
    Accessor functions for heapInitialised and placementAddress.
  **/
  bool isHeapInitialised() { return heapInitialised; }
  Address getPlacementAddress() { return placementAddress; }
  void setPlacementAddress(Address a) { placementAddress = a; }
  
  /**
    Forces the placementAddress variable to be PAGE_SIZE aligned.
  **/
  void alignPlacementAddress();
  
  /**
    Changes the original stack given by the bootloader to one at
    a virtual memory location defined at compile time.
  **/
  void remapStack();
  
  /**
    Accessor functions for kernelDirectory, currentDirectory
  **/
  PageDirectory *getKernelDirectory()  { return kernelDirectory; }
  PageDirectory *getCurrentDirectory() { return currentDirectory; }
  void setCurrentDirectory(PageDirectory *p) { currentDirectory = p; }

  /**
    Finds a free frame (swaps out if necessary) and allocates it to p.
  **/
  void allocFrame(Page *p, bool isKernel=true, bool isWriteable=true);
  /**
    Finds a free frame and returns it.
  **/
  Address allocFrame();

  /**
    Removes the frame under p's control and returns it to the pool.
  **/
  void freeFrame(Page *p);
  /**
    Adds the previously allocated frame 'frame' and returns it to the pool.
  **/
  void freeFrame(Address frame);
  
  /**
    Causes the given range of virtual memory to get allocated physical
    memory.
  **/
  void allocateRange(Address startAddress, Address size);

  /**
    Returns the size of the kernel heap. For analysis purposes.
  **/
  u32int getKernelHeapSize();
  u32int getUserHeapSize();
  
  /**
    Checks the kernel and user heap for integrity.
  **/
  void checkIntegrity();
  
private:

  /**
    Array of frames to describe physical memory state.
  **/
  BitArray *frames;
  
  /**
    Tells us how many frames there are.
  **/
  u32int nFrames;

  /**
    Has the kernel heap been initialised yet?
  **/
  bool heapInitialised;
  
  /**
    The kernel heap
  **/
  Heap heap;
  
  /**
    The user-mode shared heap
  **/
  Heap userHeap;
  
  /**
    Before the heap is initialised, this holds the next available location
    for 'placement new' to be called on.
  **/
  Address placementAddress;
  
  /**
    The currently active page directory
  **/
  PageDirectory *currentDirectory;
  
  /**
    Pointer to the "master" page directory. This holds page table pointers for kernel
    space. All other page directories must match the entries in here to maintain easy
    consistency of kernel-space over memory spaces.
  **/
  PageDirectory *kernelDirectory;
};

#ifdef LANG_X86
# include "../arch/x86/MemoryManagement.h"
#endif // LANG_X86

#endif
